zh_yogurt_category = [
    '酸奶#整体', '酸奶#味道', '酸奶#香味', '酸奶#口感', '酸奶#色泽', '酸奶#新鲜度', '酸奶#包装'
]

# 情感极性 polarity
english_sentiment = ['negative', 'neutral', 'positive']


def get_jsonl(data_path):
    """
    将data中的指定目录下的train.jsonl文件读取下来（get_category_sentiment_num_list中处理），并按照标签转为字典形式
    :param data_path: train dev test jsonl数据存放路径
    :return:
    """
    with open(data_path, 'r', encoding='utf-8') as fp:
        lines = fp.readlines()

    data_list = []
    # count = 0
    for line in lines:
        line = eval(line)
        # if len(line['sentence']) > 200:
        #     count += 1
        #     continue
        data_list.append(line)
    return data_list


# return [[c1, c2, ...], {c1: id1, c2: id2, ...}, {id1: c1, id2: c2, ...}]
def get_aspect_category(task: str, data_type: str):
    categorys = zh_yogurt_category

    # 字典生成式 enumerate 从0开始id递增
    category2id = {c: i for i, c in enumerate(categorys)}
    id2category = {i: c for i, c in enumerate(categorys)}
    return [categorys] + [category2id] + [id2category]


# return [[sentiment], {sentiment: id}, {id: sentiment}]
def get_sentiment(task):
    sentiment = english_sentiment
    sentiment2id = {s: i for i, s in enumerate(sentiment)}
    id2sentiment = {i: s for i, s in enumerate(sentiment)}
    return [sentiment] + [sentiment2id] + [id2sentiment]


# return [category_num_list, sentiment_num_list] 训练集中的[方面类别]和[情感极性]数量统计
def get_category_sentiment_num_list(args):
    task, data_type, data_path = args.task, args.data_type, args.data_path
    data_path += 'train.jsonl'
    categorys = zh_yogurt_category
    sentiment = english_sentiment

    category_dict = {i: 0 for i in categorys}
    sentiment_dict = {i: 0 for i in sentiment}
    data_list = get_jsonl(data_path)
    count = 0
    for d in data_list:
        labels = d['labels']
        for label in labels:
            cate, senti = label[1], label[-1]
            category_dict[cate] += 1
            sentiment_dict[senti] += 1
            count += 1

    category_num_list, sentiment_num_list = [0 for _ in range(len(category_dict))], [0 for _ in
                                                                                     range(len(sentiment_dict))]
    for i, key in enumerate(category_dict):
        category_num_list[i] = category_dict[key]
    for i, key in enumerate(sentiment_dict):
        sentiment_num_list[i] = sentiment_dict[key]

    return [category_num_list, sentiment_num_list]
